package org.java.vm;

import org.java.vm.command.ArithmeticHandler;
import org.java.vm.command.*;

import java.io.File;
import java.util.*;

public class CodeWriter {

    private final CommandWriter commandWriter;
    private final ArithmeticHandler arithmeticHandler;
    private final List<ArithmeticCommand> commandHandlers = new ArrayList<>();
    private String currentFunction;
    private String variablePrefix;

    public CodeWriter(File asmFile) {
        arithmeticHandler = new ArithmeticHandler();
        commandWriter = new CommandWriter(asmFile);
        initArithmeticCommand();
    }

    private void initArithmeticCommand() {
        commandHandlers.add(new Add(arithmeticHandler));
        commandHandlers.add(new Sub(arithmeticHandler));
        commandHandlers.add(new Neg(arithmeticHandler));
        commandHandlers.add(new Eq(arithmeticHandler));
        commandHandlers.add(new Lt(arithmeticHandler));
        commandHandlers.add(new Gt(arithmeticHandler));
        commandHandlers.add(new And(arithmeticHandler));
        commandHandlers.add(new Or(arithmeticHandler));
        commandHandlers.add(new Not(arithmeticHandler));
    }

    public void writeInit() {
        commandWriter.output("// init");
        commandWriter.output("@256", "D=A", "@SP", "M=D");
        writeCall("Sys.init", 0);
    }

    public void setFileName(String vmFileName) {
        this.variablePrefix = vmFileName.replace("vm", "");
        this.arithmeticHandler.setVariablePrefix(variablePrefix);
    }

    public void writeArithmetic(String arithmetic) {
        commandWriter.output("// " + arithmetic);

//        operateTwoDigit(arithmetic);

        List<String> asmCommands = new ArrayList<>();
        for (ArithmeticCommand commandHandler : commandHandlers) {
            if (commandHandler.canCalculate(arithmetic)) {
                asmCommands = commandHandler.calculate();
            }
        }
        commandWriter.output(asmCommands.toArray(new String[0]));
    }

    private void operateTwoDigit(String arithmetic) {
        writePop2Register("R14");
        writePop2Register("R13");
        List<String> commands = new ArrayList<>();
        commands.add("@R14");
        commands.add("D=M");
        commands.add("@R13");
        switch (arithmetic) {
            case "add":
                commands.add("M=D+M");
                break;
            case "sub":
                commands.add("M=M-D");
            case "and":
                commands.add("M=D&M");
            case "or":
                commands.add("M=D|M");
            case "eq":
                commands.add("D=M-D");
                commands.add("@");
        }
        commandWriter.output(commands.toArray(new String[0]));
        pushRegister("R13");
    }

    public void writePushPop(CommandType commandType, String arg1, int arg2) {
        commandWriter.output("//" + commandType.name() + " " + arg1 + " " + arg2);
        List<String> commands = new ArrayList<>();
        switch (commandType) {
            case C_PUSH:
                push(arg1, arg2);
                break;
            case C_POP:
                commands = arithmeticHandler.pop(arg1, arg2);
        }
        commandWriter.output(commands.toArray(new String[0]));
    }

    private void push(String arg1, int arg2) {
        if (arg1.equalsIgnoreCase("constant")) {
            pushConstant(arg2 + "");
        } else if (arg1.equalsIgnoreCase("static")) {
            pushRegister(variablePrefix + arg2);
        } else if (arg1.equalsIgnoreCase("pointer")) {
            pushRegisterData("R3", arg2, false);
        } else if (arg1.equalsIgnoreCase("temp")) {
            pushRegisterData("R5", arg2, false);
        } else if (arg1.equalsIgnoreCase("argument")){
            pushRegisterData("ARG", arg2, true);
        } else if (arg1.equalsIgnoreCase("local")) {
            pushRegisterData("LCL", arg2, true);
        } else if (arg1.equalsIgnoreCase("this")) {
            pushRegisterData("THIS", arg2, true);
        } else if (arg1.equalsIgnoreCase("that")) {
            pushRegisterData("THAT", arg2, true);
        }
    }

    public void writeLabel(String label) {
        commandWriter.output("// label " + label);
        if (currentFunction != null) {
            label = currentFunction + "$" + label;
        }
        commandWriter.output("(" + label + ")");
    }

    public void writeIf(String label) {
        commandWriter.output("// if-goto " + label);
        if (currentFunction != null) {
            label = currentFunction + "$" + label;
        }
        // 检查栈上数据是否大于0
        popStack();
        commandWriter.output("@" + label, "D;JNE");
    }

    private void popStack() {
        commandWriter.output("@SP", "M=M-1", "@SP", "A=M", "D=M");
    }

    public void writeGoto(String label) {
        commandWriter.output("// goto " + label);
        if (currentFunction != null) {
            label = currentFunction + "$" + label;
        }
        commandWriter.output("@" + label, "0;JMP");
    }

    public void writeFunction(String functionName, Integer numLocals) {
        commandWriter.output("// function " + functionName + " " + numLocals);
        commandWriter.output("(" + functionName + ")");
        for (int i = 0; i < numLocals; i++) {
            // LCL进栈，初始化为0
            commandWriter.output("@SP", "A=M", "M=0", "@SP", "M=M+1");
        }

        currentFunction = functionName;
    }

    public void writeReturn() {
        commandWriter.output("// return");
        // LCL记录到R13  FRAME = LCL
        copyRegister("LCL", "R13");
        // RET = *(FRAME-5)  RET存入 R14
        restoreRegisterByBaseRegisterOffsetMemory("R14", "R13", 5);
        // *ARG = pop();
        pop2Segment("ARG", 0);
        // 恢复调用者的SP SP = ARG + 1
        commandWriter.output("@ARG", "D=M", "@SP", "M=D+1");
        // 恢复THAT   THAT = *(FRAME - 1)
        restoreRegisterByBaseRegisterOffsetMemory("THAT", "R13", 1);
        // 恢复THIS    THIS = *(FRAME - 2)
        restoreRegisterByBaseRegisterOffsetMemory("THIS", "R13", 2);
        // 恢复ARG
        restoreRegisterByBaseRegisterOffsetMemory("ARG", "R13", 3);
        // 恢复LCL
        restoreRegisterByBaseRegisterOffsetMemory("LCL", "R13", 4);
        // goto RET
        commandWriter.output("@R14", "A=M", "0;JMP");
    }

    /**
     * register = * (baseRegister - index)
     *
     * @param register
     * @param baseRegister
     * @param index
     * @return
     */
    private void restoreRegisterByBaseRegisterOffsetMemory(String register, String baseRegister, int index) {
        List<String> commands = new ArrayList<>();
        commands.add("@" + index);
        commands.add("D=A");
        commands.add("@" + baseRegister);
        commands.add("A=M-D");
        commands.add("D=M");
        commands.add("@" + register);
        commands.add("M=D");

        commandWriter.output(commands.toArray(new String[0]));
    }

    /**
     * register = SP - index
     *
     * @param register
     * @param index
     * @return
     */
    private void setRegisterBySP(String register, int index) {
        List<String> commands = new ArrayList<>();
        commands.add("@" + index);
        commands.add("D=A");
        commands.add("@SP");
        commands.add("D=M-D");
        commands.add("@" + register);
        commands.add("M=D");

        commandWriter.output(commands.toArray(new String[0]));
    }

    private void copyRegister(String from, String to) {
        List<String> commands = new ArrayList<>();
        commands.add("@" + from);
        commands.add("D=M");
        commands.add("@" + to);
        commands.add("M=D");
        commandWriter.output(commands.toArray(new String[0]));
    }

    private void writePop2Register(String register) {
        List<String> commands = new ArrayList<>();
        commands.add("@SP");
        commands.add("M=M-1");
        commands.add("@SP");
        commands.add("A=M");
        commands.add("D=M");
        commands.add("@" + register);
        commands.add("M=D");
        commandWriter.output(commands.toArray(new String[0]));
    }

    public void close() {
        commandWriter.close();
    }

    public void writeCall(String functionName, Integer numArgs) {
        commandWriter.output("// call " + functionName + " " + numArgs);

        // push return-address
        String returnAddress = IDGenerator.uniqueID("return_address_");
        pushConstant(returnAddress);
        // push local
        pushRegister("LCL");
        // push arg
        pushRegister("ARG");
        // push this
        pushRegister("THIS");
        // push that
        pushRegister("THAT");
        // restore arg ARG = SP -n - 5
        setRegisterBySP("ARG", numArgs + 5);
        // restore lcl
        setRegisterBySP("LCL", 0);
        // goto function
        commandWriter.output("@" + functionName, "0;JMP");
        // return-address label
        commandWriter.output("(" + returnAddress + ")");
    }

    /**
     * *(arg1 + index) = * (sp--)
     *
     * @param segment
     * @param index
     */
    private void pop2Segment(String segment, int index) {
        // D = segment + index
        commandWriter.output("@" + index, "D=A", "@" + segment, "D=D+M");
        // *sp = D
        commandWriter.output("@SP", "A=M", "M=D");
        // D = *(sp - 1)
        commandWriter.output("@SP", "A=M-1", "D=M");
        // *(segment + index) = D
        commandWriter.output("@SP", "A=M", "A=M", "M=D");
        // sp - 1
        commandWriter.output("@SP", "M=M-1");
    }

    private void pushConstant(String label) {
        List<String> commands = new ArrayList<>();
        // push return-address
        commands.add("@" + label);
        commands.add("D=A");
        commands.add("@SP");
        commands.add("A=M");
        commands.add("M=D");
        commands.add("@SP");
        commands.add("M=M+1");
        commandWriter.output(commands.toArray(new String[0]));
    }


    private void pushRegister(String register) {
        List<String> commands = new ArrayList<>();
//         push return-address
        commands.add("@" + register);
        commands.add("D=M");
        commands.add("@SP");
        commands.add("A=M");
        commands.add("M=D");
        commands.add("@SP");
        commands.add("M=M+1");
        commandWriter.output(commands.toArray(new String[0]));
    }

    private void pushRegisterData(String register, int offset, boolean isRegisterMemory) {
        List<String> commands = new ArrayList<>();
        commands.add("@" + offset);
        commands.add("D=A");
        commands.add("@" + register);
        if (isRegisterMemory) {
            commands.add("A=D+M");
        } else {
            commands.add("A=D+A");
        }
        commands.add("D=M");

        commands.add("@SP");
        commands.add("A=M");
        commands.add("M=D");

        commands.add("@SP");
        commands.add("M=M+1");
        commandWriter.output(commands.toArray(new String[0]));
    }
}
